// File: jmcomic-core/src/main/java/dev/jukomu/common/base/concurrent/MultiTaskExecutor.java
package dev.jukomu.common.base.concurrent;

import dev.jukomu.common.util.args.ArgsUtils;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.concurrent.*;
import java.util.function.Consumer;
import java.util.function.Function;

/**
 * 对应Python的MultiTaskLauncher，用于管理多线程任务的启动和等待。
 */
public class MultiTaskExecutor {

    // 默认的StopFlag实例，可以在创建Executor时传入自定义的
    public static final StopFlag DEFAULT_STOP_FLAG = new StopFlag("multi_task_stop_flag");

    private final Map<String, Object> taskMetadata; // 任务元数据，对应Python的task_metadata
    private final StopFlag stopFlag; // 停止标志
    private final List<Future<?>> taskFutures; // 存储所有提交的任务Future
    private final ExecutorService executorService; // 线程池

    /**
     * 构造函数。
     *
     * @param taskMetadata   任务元数据，用于配置线程池（例如线程名）。
     * @param stopFlag       停止标志，如果为null则使用默认的StopFlag。
     * @param threadPoolSize 线程池大小。如果为0或负数，则使用CachedThreadPool；否则使用FixedThreadPool。
     */
    public MultiTaskExecutor(Map<String, Object> taskMetadata, StopFlag stopFlag, int threadPoolSize) {
        this.taskMetadata = taskMetadata != null ? new ConcurrentHashMap<>(taskMetadata) : new ConcurrentHashMap<>();
        this.stopFlag = stopFlag != null ? stopFlag : DEFAULT_STOP_FLAG;
        this.taskFutures = new ArrayList<>();

        if (threadPoolSize <= 0) {
            this.executorService = Executors.newCachedThreadPool(); // 对应Python的daemon=True的动态线程池
        } else {
            this.executorService = Executors.newFixedThreadPool(threadPoolSize); // 固定大小线程池
        }
    }

    /**
     * 构造函数，使用默认线程池大小（根据CPU核心数）。
     *
     * @param taskMetadata 任务元数据。
     * @param stopFlag     停止标志。
     */
    public MultiTaskExecutor(Map<String, Object> taskMetadata, StopFlag stopFlag) {
        this(taskMetadata, stopFlag, Runtime.getRuntime().availableProcessors()); // 默认线程池大小为CPU核心数
    }

    /**
     * 构造函数，使用默认任务元数据和停止标志。
     */
    public MultiTaskExecutor() {
        this(null, null, 0); // 默认CachedThreadPool
    }

    /**
     * 检查停止标志，如果应该停止则抛出InterruptedException。
     *
     * @throws InterruptedException 如果应该停止。
     */
    private void checkStopFlag() throws InterruptedException {
        stopFlag.ifStopRaise();
    }

    /**
     * 创建并提交一个任务。
     *
     * @param target 任务执行逻辑（Runnable或Callable）。
     * @param args   可变参数列表。
     * @param kwargs 关键字参数Map。
     * @param <T>    Callable的返回类型。
     * @return 任务的Future对象。
     */
    public <T> Future<T> createTask(Callable<T> target, Object[] args, Map<String, Object> kwargs) {
        Callable<T> wrappedCallable = () -> {
            stopFlag.markRun(); // 标记当前线程为运行状态
            try {
                return target.call();
            } catch (Exception e) {
                // 将异常重新包装以便在Future.get()时抛出
                throw new ExecutionException("Task execution failed", e);
            }
        };
        Future<T> future = executorService.submit(wrappedCallable);
        taskFutures.add(future);
        return future;
    }

    /**
     * 创建并提交一个Runnable任务。
     *
     * @param target 任务执行逻辑（Runnable）。
     * @param args   可变参数列表。
     * @param kwargs 关键字参数Map。
     * @return 任务的Future对象。
     */
    public Future<?> createTask(Runnable target, Object[] args, Map<String, Object> kwargs) {
        Runnable wrappedRunnable = () -> {
            stopFlag.markRun(); // 标记当前线程为运行状态
            target.run();
        };
        Future<?> future = executorService.submit(wrappedRunnable);
        taskFutures.add(future);
        return future;
    }

    /**
     * 等待所有已提交的任务完成。
     * 如果在等待期间收到停止信号，则抛出InterruptedException。
     *
     * @throws InterruptedException 如果等待被中断或收到停止信号。
     */
    public void waitFinish() throws InterruptedException {
        try {
            for (Future<?> future : taskFutures) {
                while (!future.isDone()) {
                    checkStopFlag(); // 每次检查时都检查停止标志
                    try {
                        future.get(100, TimeUnit.MILLISECONDS); // 尝试等待100ms
                    } catch (java.util.concurrent.TimeoutException ignored) {
                        // Timeout, continue checking stopFlag
                    } catch (ExecutionException e) {
                        // If task itself threw an exception, unwrap and rethrow or log
                        Throwable cause = e.getCause();
                        if (cause instanceof InterruptedException) {
                            throw (InterruptedException) cause;
                        } else if (cause instanceof RuntimeException) {
                            throw (RuntimeException) cause;
                        } else {
                            throw new RuntimeException("Task completed with unexpected exception", cause);
                        }
                    }
                }
            }
        } finally {
            shutdownExecutor();
        }
    }

    /**
     * 关闭ExecutorService，等待任务完成。
     */
    private void shutdownExecutor() {
        executorService.shutdown(); // Initiates an orderly shutdown
        try {
            if (!executorService.awaitTermination(60, TimeUnit.SECONDS)) { // Wait up to 60 seconds for tasks to finish
                executorService.shutdownNow(); // Force shutdown
                if (!executorService.awaitTermination(60, TimeUnit.SECONDS)) {
                    System.err.println("ExecutorService did not terminate.");
                }
            }
        } catch (InterruptedException ie) {
            executorService.shutdownNow();
            Thread.currentThread().interrupt();
        }
    }

    /**
     * 暂停当前线程一段时间，并在暂停期间检查停止标志。
     *
     * @param durationSeconds 暂停时长（秒）。
     * @throws InterruptedException 如果暂停被中断或收到停止信号。
     */
    public void pause(double durationSeconds) throws InterruptedException {
        stopFlag.sleepOrRaise(durationSeconds);
    }

    /**
     * 获取已提交任务的数量。
     *
     * @return 任务数量。
     */
    public int size() {
        return taskFutures.size();
    }

    /**
     * 创建一个守护线程Executor。
     * 在Java中，所有非守护线程退出后，守护线程会自动退出。
     *
     * @return MultiTaskExecutor实例。
     */
    public static MultiTaskExecutor buildDaemon() {
        // In Java, there's no direct "daemon" flag for ExecutorService in the same way as Python Thread.
        // Threads created by Executors.newCachedThreadPool() are often daemon by default in common JDKs,
        // but it's not guaranteed by ExecutorService API.
        // For explicit daemon behavior, custom ThreadFactory is needed.
        // For simplicity, we use newCachedThreadPool and rely on JVM's behavior.
        // If strict daemon behavior is required, ThreadFactory should be added.
        return new MultiTaskExecutor(null, null, 0); // CachedThreadPool for daemon-like behavior
    }

    // --- 对应Python的顶层函数 ---

    /**
     * 对应Python的multi_task_launcher和multi_thread_launcher。
     * 批量启动任务，使用线程池。
     *
     * @param iterObjs         可迭代对象，每个元素将被传递给applyEachObjFunc。
     * @param applyEachObjFunc 对每个元素执行的函数。
     * @param waitFinish       是否等待所有任务完成。
     * @param <T>              iterObjs中元素的类型。
     * @return MultiTaskExecutor实例。
     * @throws InterruptedException 如果任务执行被中断。
     */
    public static <T> MultiTaskExecutor launch(Collection<T> iterObjs, Consumer<T> applyEachObjFunc, boolean waitFinish) throws InterruptedException {
        MultiTaskExecutor executor = new MultiTaskExecutor();
        for (T obj : iterObjs) {
            executor.createTask(() -> {
                applyEachObjFunc.accept(obj);
                return null;
            }, null, null);
        }
        if (waitFinish) {
            executor.waitFinish();
        }
        return executor;
    }

    /**
     * 对应Python的thread_pool_executor。
     * 使用固定大小的线程池处理任务。
     *
     * @param iterObjs         可迭代对象，每个元素将被传递给applyEachObjFunc。
     * @param applyEachObjFunc 对每个元素执行的函数。
     * @param waitFinish       是否等待所有任务完成。
     * @param maxWorkers       线程池最大工作线程数。
     * @param stopFlag         停止标志。
     * @param <T>              iterObjs中元素的类型。
     * @return MultiTaskExecutor实例。
     * @throws InterruptedException 如果任务执行被中断。
     */
    public static <T> MultiTaskExecutor threadPoolExecutor(Collection<T> iterObjs, Consumer<T> applyEachObjFunc, boolean waitFinish, int maxWorkers, StopFlag stopFlag) throws InterruptedException {
        MultiTaskExecutor executor = new MultiTaskExecutor(null, stopFlag, maxWorkers);
        for (T obj : iterObjs) {
            executor.createTask(() -> {
                applyEachObjFunc.accept(obj);
                return null;
            }, null, null);
        }
        if (waitFinish) {
            executor.waitFinish();
        }
        return executor;
    }


    /**
     * 对应Python的multi_call。
     * 批量调用函数并收集结果。
     *
     * @param func     要调用的函数。
     * @param iterObjs 可迭代对象，每个元素将作为参数传递给func。
     * @param <T>      iterObjs中元素的类型。
     * @param <R>      func的返回类型。
     * @return 包含输入对象及其对应结果的Map。
     * @throws InterruptedException 如果任务执行被中断。
     */
    public static <T, R> Map<T, R> multiCall(Function<T, R> func, Collection<T> iterObjs) throws InterruptedException {
        MultiTaskExecutor executor = new MultiTaskExecutor();
        Map<T, Future<R>> futureMap = new ConcurrentHashMap<>();

        for (T obj : iterObjs) {
            futureMap.put(obj, executor.createTask(() -> func.apply(obj), null, null));
        }

        executor.waitFinish(); // Wait for all tasks to complete

        Map<T, R> resultMap = new ConcurrentHashMap<>();
        for (Map.Entry<T, Future<R>> entry : futureMap.entrySet()) {
            try {
                resultMap.put(entry.getKey(), entry.getValue().get());
            } catch (ExecutionException e) {
                // Log or rethrow the cause of the exception if needed
                throw new RuntimeException("Failed to get result for key: " + entry.getKey(), e.getCause());
            }
        }
        return resultMap;
    }

    /**
     * 对应Python的invoke_all。
     * 批量调用带有不同参数和函数的任务。
     *
     * @param argsFuncList 包含参数和函数的列表，每个元素是一个Map.Entry，key是参数，value是函数。
     *                     参数可以是单个值、List或Map。
     * @param wait         是否等待所有任务完成。
     * @return 所有任务的Future列表，如果wait为true，则返回结果列表。
     * @throws InterruptedException 如果任务执行被中断。
     */
    public static <R> List<Future<R>> invokeAll(List<Map.Entry<Object, Function<Object[], R>>> argsFuncList, boolean wait) throws InterruptedException {
        MultiTaskExecutor executor = new MultiTaskExecutor();
        List<Future<R>> futureList = new ArrayList<>();

        for (Map.Entry<Object, Function<Object[], R>> entry : argsFuncList) {
            Object arg = entry.getKey();
            Function<Object[], R> func = entry.getValue();

            // Use ArgsUtils to process args if it's a list/map, otherwise wrap it
            Object[] processedArgs = ArgsUtils.argToTuple(arg);

            futureList.add(executor.createTask(() -> func.apply(processedArgs), processedArgs, null));
        }

        if (wait) {
            executor.waitFinish();
            // Get results if wait is true
            List<R> results = new ArrayList<>();
            for (Future<R> future : futureList) {
                try {
                    results.add(future.get());
                } catch (ExecutionException e) {
                    throw new RuntimeException("Failed to get result from future", e.getCause());
                }
            }
            // Returning a List<R> for wait=true is more intuitive, but Python returns Future list.
            // Sticking to Python's behavior, we always return List<Future<R>> and let consumer get results.
            // If the original Python behavior is to return results for wait=true,
            // the return type here needs to be adjusted.
            // For now, consistent return type of List<Future<R>>.
        }
        return futureList;
    }

    /**
     * 对应Python的CacheRunner。
     * 这是一个简单的包装，CacheRunner在Java中更适合作为Executor内部的缓存Future。
     *
     * @param func 要缓存执行的函数。
     * @param <R>  函数的返回类型。
     * @return 一个Future，其结果是func的执行结果。
     */
    public static <R> Future<R> cacheRun(Callable<R> func) {
        // In Java, CacheRunner is typically implemented by submitting a task to an ExecutorService
        // and holding onto the Future object. The Future itself acts as the cache.
        // We'll create a single-task executor for this specific run.
        ExecutorService singleExecutor = Executors.newSingleThreadExecutor();
        Future<R> future = singleExecutor.submit(() -> {
            try {
                return func.call();
            } finally {
                singleExecutor.shutdown(); // Ensure the single thread executor shuts down
            }
        });
        return future;
    }
}